import chisel3._
import chisel3.util._
import Control._

class IDU(xlen: Int) extends Module {
  val io = IO(new Bundle {
    val in  = Flipped(Decoupled(new MessageFD(xlen)))
    val out = Decoupled(new MessageDE(xlen))

    // WBU => IDU
    val in_wb = Flipped(new MessageWD(xlen))

    // MEU => IDU
    val meu_wb_en      = Input(Bool())
    val meu_wb_sel     = Input(UInt(2.W))
    val meu_wb_Alu_out = Input(UInt(xlen.W))
    val meu_rd_data    = Input(UInt(xlen.W))
    val meu_wb_addr    = Input(UInt(5.W))
    val meu_out_valid  = Input(Bool())
    val flush          = Input(Bool())

    // EXU => IDU
    val exu_wb_en     = Input(Bool())
    val exu_wb_sel    = Input(UInt(2.W))
    val exu_Alu_out   = Input(UInt(xlen.W))
    val exu_wb_addr   = Input(UInt(5.W))
    val exu_out_valid = Input(Bool())
  })

  // pipeline registers (ifu => idu)
  val idu_inst     = RegInit(0.U(xlen.W))
  val idu_pc       = RegInit(0.U(xlen.W))
  val idu_wb_addr  = RegInit(0.U(5.W))
  val idu_rs1_data = RegInit(0.U(xlen.W))
  val idu_rs2_data = RegInit(0.U(xlen.W))

  // valid for output
  val valid = RegInit(0.B)

  // ready for output
  val ready = RegInit(1.B)

  val decoder = Module(new Decoder(xlen))
  val rf = Module(new RegisterFile(xlen, xlen))

  idu_inst     := Mux(io.flush, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.inst,     idu_inst))
  idu_pc       := Mux(io.flush, 0.U, Mux(io.in.valid && io.in.ready, io.in.bits.pc,       idu_pc))

  decoder.io.inst := idu_inst

  io.out.bits.inst    := idu_inst
  io.out.bits.pc      := idu_pc
  io.out.bits.PC_sel  := decoder.io.PC_sel
  io.out.bits.Imm_sel := decoder.io.Imm_sel
  io.out.bits.Alu_op  := decoder.io.Alu_op
  io.out.bits.st_type := decoder.io.st_type
  io.out.bits.ld_type := decoder.io.ld_type
  io.out.bits.br_type := decoder.io.br_type
  io.out.bits.A_sel   := decoder.io.A_sel
  io.out.bits.B_sel   := decoder.io.B_sel
  io.out.bits.wb_sel  := decoder.io.wb_sel
  io.out.bits.csr_cmd := decoder.io.csr_cmd
  io.out.bits.wb_en   := decoder.io.wb_en

  io.out.bits.wb_addr  := idu_inst(11, 7)
  // io.out.bits.rs1_data := rf.io.rs1_data
  // io.out.bits.rs2_data := rf.io.rs2_data

  // hazard
  val rs1hazard_exu = io.exu_out_valid && io.exu_wb_en && (io.exu_wb_addr === idu_inst(19, 15))
  val rs2hazard_exu = io.exu_out_valid && io.exu_wb_en && (io.exu_wb_addr === idu_inst(24, 20))
  val rs1hazard_meu = io.meu_out_valid && io.meu_wb_en && (io.meu_wb_addr === idu_inst(19, 15))
  val rs2hazard_meu = io.meu_out_valid && io.meu_wb_en && (io.meu_wb_addr === idu_inst(24, 20))
  val rs1hazard_wbu = io.in_wb.valid   && io.in_wb.wb_en && (io.in_wb.wb_addr === idu_inst(19, 15))
  val rs2hazard_wbu = io.in_wb.valid   && io.in_wb.wb_en && (io.in_wb.wb_addr === idu_inst(24, 20))
  val isRAW = rs1hazard_exu | rs2hazard_exu | rs1hazard_meu | rs2hazard_meu | rs1hazard_wbu | rs2hazard_wbu

  io.out.bits.rs1_data := Mux(idu_inst(19, 15) === 0.U, 0.U, Mux(!io.in_wb.valid || idu_inst(19, 15) =/= io.in_wb.wb_addr, rf.io.rs1_data, MuxLookup(io.in_wb.wb_sel, io.in_wb.Alu_out)(
    Seq(WB_SEL_PC_4 -> (io.in_wb.pc + 4.U),
        WB_SEL_ALU  -> io.in_wb.Alu_out,
        WB_SEL_MEM  -> io.in_wb.rd_data,
        WB_SEL_CSR  -> io.in_wb.csr_out
      ))))
  io.out.bits.rs2_data := Mux(idu_inst(24, 20) === 0.U, 0.U, Mux(!io.in_wb.valid || idu_inst(24, 20) =/= io.in_wb.wb_addr, rf.io.rs2_data, MuxLookup(io.in_wb.wb_sel, io.in_wb.Alu_out)(
    Seq(WB_SEL_PC_4 -> (io.in_wb.pc + 4.U),
        WB_SEL_ALU  -> io.in_wb.Alu_out,
        WB_SEL_MEM  -> io.in_wb.rd_data,
        WB_SEL_CSR  -> io.in_wb.csr_out
      ))))

  rf.io.wb_en    := io.in_wb.valid & io.in_wb.wb_en
  rf.io.Alu_out  := io.in_wb.Alu_out
  rf.io.csr_out  := io.in_wb.csr_out
  rf.io.pc_4     := io.in_wb.pc + 4.U
  rf.io.wb_sel   := io.in_wb.wb_sel
  rf.io.wb_addr  := io.in_wb.wb_addr
  rf.io.rs1_addr := idu_inst(19, 15)
  rf.io.rs2_addr := idu_inst(24, 20)
  rf.io.rd_data  := io.in_wb.rd_data

  valid := Mux(io.out.valid & io.out.ready, 0.B, Mux(io.in.valid & io.in.ready, 1.B, valid))
  io.out.valid := valid && Mux(!isRAW, 1.B, Mux(io.in_wb.valid && io.in_wb.wb_en && (io.in_wb.wb_addr === idu_inst(19, 15) || io.in_wb.wb_addr === idu_inst(24, 20)), 1.B, 0.B))

  ready := Mux(io.in.valid & io.in.ready, 0.B, Mux(io.out.valid & io.out.ready, 1.B, ready))
  io.in.ready  := ready && ~io.flush
}
